# -*- coding: utf-8 -*-
"""
Tavily搜索服务 - 专为AI优化的搜索API
"""

import os
import asyncio
import logging
import aiohttp
import json
from typing import List, Dict, Optional
from pathlib import Path

logger = logging.getLogger(__name__)

class SearchResult:
    def __init__(self, title: str, url: str, snippet: str):
        self.title = title
        self.url = url
        self.snippet = snippet

class TavilySearchService:
    def __init__(self):
        self.timeout = 10.0  # Tavily搜索超时时间
        self.max_results = 5  # 最大搜索结果数
        self.token_limit = 500  # 搜索结果token限制
        self.api_url = "https://api.tavily.com/search"
        
        # 从环境变量或.env文件获取API密钥
        self.api_key = self._load_api_key()
        if not self.api_key:
            logger.warning("TAVILY_API_KEY 环境变量未设置，搜索功能将不可用")
        else:
            logger.info(f"Tavily搜索服务初始化成功，API密钥: {self.api_key[:10]}...")
    
    def _load_api_key(self) -> Optional[str]:
        """
        从环境变量或.env文件加载API密钥
        """
        # 首先尝试从环境变量获取
        api_key = os.getenv('TAVILY_API_KEY')
        if api_key:
            return api_key
        
        # 如果环境变量不存在，尝试从.env文件读取
        env_file = Path(__file__).parent / '.env'
        if env_file.exists():
            try:
                with open(env_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        line = line.strip()
                        if line.startswith('TAVILY_API_KEY='):
                            return line.split('=', 1)[1]
            except Exception as e:
                logger.error(f"读取.env文件失败: {str(e)}")
        
        return None
    
    async def search(self, query: str) -> List[SearchResult]:
        """
        执行Tavily搜索
        
        Args:
            query: 搜索查询
            
        Returns:
            List[SearchResult]: 搜索结果列表
        """
        if not self.api_key:
            logger.error("Tavily API密钥未设置，无法执行搜索")
            return []
        
        try:
            logger.info(f"执行Tavily搜索: 查询='{query}'")
            
            # 构建请求数据
            payload = {
                "api_key": self.api_key,
                "query": query,
                "search_depth": "basic",
                "max_results": self.max_results,
                "include_answer": True,
                "include_raw_content": False
            }
            
            # 使用aiohttp进行异步HTTP请求
            timeout = aiohttp.ClientTimeout(total=self.timeout)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(self.api_url, json=payload) as response:
                    if response.status == 200:
                        data = await response.json()
                    else:
                        logger.error(f"Tavily API请求失败: 状态码={response.status}")
                        return []
            
            results = []
            if data and 'results' in data:
                for item in data['results'][:self.max_results]:
                    result = SearchResult(
                        title=item.get('title', '无标题'),
                        url=item.get('url', ''),
                        snippet=item.get('content', '')[:300]  # 限制长度
                    )
                    results.append(result)
                
                logger.info(f"Tavily搜索成功: 查询='{query}', 结果数={len(results)}")
            else:
                logger.warning(f"Tavily搜索返回空结果: 查询='{query}'")
            
            return results
            
        except Exception as e:
            logger.error(f"Tavily搜索失败: 查询='{query}', 错误={str(e)}")
            return []
    
    def should_trigger_search(self, query: str) -> bool:
        """
        判断是否应该触发搜索
        
        Args:
            query: 用户查询
            
        Returns:
            bool: 是否需要搜索
        """
        # 如果API密钥不可用，不触发搜索
        if not self.api_key:
            return False
        
        # 查询长度过短，可能不需要搜索
        if len(query) < 3:
            return False
        
        # 检查是否包含时间敏感词汇
        time_sensitive = [
            '最新', '新出', '最近', '今年', '2023', '2024', '2025', 
            '现在', '目前', '当前', '实时', '最新的'
        ]
        for word in time_sensitive:
            if word in query:
                return True
        
        # 检查是否包含事实性查询词汇
        factual_words = [
            '什么', '如何', '怎么', '为什么', '哪里', '谁', '何时', 
            '多少', '价格', '费用', '地址', '电话', '时间'
        ]
        for word in factual_words:
            if word in query:
                return True
        
        # 检查是否包含需要实时信息的关键词
        realtime_keywords = [
            '新闻', '股价', '天气', '汇率', '比分', '票价', 
            '开放时间', '营业时间', '招聘', '招生'
        ]
        for keyword in realtime_keywords:
            if keyword in query:
                return True
        
        return False
    
    def format_results_for_context(self, results: List[SearchResult]) -> str:
        """
        将搜索结果格式化为上下文文本
        
        Args:
            results: 搜索结果列表
            
        Returns:
            str: 格式化的上下文
        """
        if not results:
            return ""
        
        context = "## 搜索结果\n\n"
        for i, result in enumerate(results, 1):
            context += f"### 结果 {i}: {result.title}\n"
            context += f"来源: {result.url}\n"
            context += f"内容: {result.snippet}\n\n"
        
        context += "---\n请基于以上搜索结果回答用户问题，优先使用搜索结果中的信息。\n"
        return context

# 全局搜索服务实例
search_service = TavilySearchService()